import numpy as np
from collections import Counter

class Node:
    """决策树的节点类"""
    def __init__(self, gini=None, num_samples=None, num_samples_per_class=None, predicted_class=None):
        self.gini = gini  # 基尼不纯度
        self.num_samples = num_samples  # 样本数量
        self.num_samples_per_class = num_samples_per_class  # 每个类别的样本数量
        self.predicted_class = predicted_class  # 预测类别
        self.feature_index = None  # 分割特征的索引
        self.threshold = None  # 分割阈值
        self.left = None  # 左子树
        self.right = None  # 右子树

class CARTDecisionTree:

    def __init__(self, max_depth=None, min_samples_split=2):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.tree = None

    def fit(self, X, y, sample_weight=None):
        # 如果 sample_weight 为空，初始化为均匀分布
        if sample_weight is None:
            sample_weight = np.ones(len(y)) / len(y)
    
        # 在划分时考虑权重
        self.tree = self.grow_tree(X, y, sample_weight)

    # 使用模型进行预测
    def predict(self, X):
        return np.array([self.single_predict(inputs) for inputs in X])

    # 计算 gini 值
    def calc_gini(self, y):
        counter = Counter(y)
        num_of_samples = len(y)
        gini_value = 1 - sum((count/num_of_samples) ** 2 for count in counter.values())
        return gini_value

    # 根据特征和阈值分割数据集
    def split_data(self, X, y, feature_index, threshold):
        left_mask = X[:, feature_index] <= threshold
        right_mask = X[:, feature_index] > threshold
        return X[left_mask], X[right_mask], y[left_mask], y[right_mask]

    # 寻找最佳划分的特征以及该特征进行划分的阈值
    def best_split(self, X, y):
        m, n = X.shape
        if m <= 1:
            return None, None

        # 计算当前节点的基尼不纯度
        num_parent = [np.sum(y == c) for c in np.unique(y)]
        best_gini = 1 - sum((num/m) ** 2 for num in num_parent)
        best_idx, best_thr = None, None
        # 遍历所有特征
        for  feature_index in range(n):
            thresholds, classes = zip(*sorted(zip(X[:, feature_index], y)))
            # 左子树各类别的数量记作一个列表，每个元素初始化为0
            num_left = [0] * len(np.unique(y))
            # 右字数各类别的数量记作一个李彪，每个元素初始化为该节点的数量
            num_right = num_parent.copy()
            for i in range(1, m):
                c = classes[i - 1]
                num_left[c] += 1
                num_right[c] -= 1
                gini_left = 1.0 - sum((num_left[x] / i) ** 2 for x in np.unique(y))
                gini_right = 1.0 - sum((num_right[x] / (m - i)) ** 2 for x in np.unique(y))
                gini = (i * gini_left + (m - i) * gini_right) / m
                
            if thresholds[i] == thresholds[i - 1]:
                    continue

            if gini < best_gini:
                best_gini = gini
                best_idx = feature_index
                best_thr = (thresholds[i] + thresholds[i - 1]) / 2

        return best_idx, best_thr
            
    def grow_tree(self, X, y, depth=0):
        """递归构建决策树"""
        num_samples_per_class = [np.sum(y == i) for i in np.unique(y)]
        predicted_class = np.argmax(num_samples_per_class)
        node = Node(
            gini=self.calc_gini(y),
            num_samples=len(y),
            num_samples_per_class=num_samples_per_class,
            predicted_class=predicted_class,
        )

        if depth < self.max_depth and len(y) >= self.min_samples_split:
            idx, thr = self.best_split(X, y)
            if idx is not None:
                X_left, X_right, y_left, y_right = self.split(X, y, idx, thr)
                node.feature_index = idx
                node.threshold = thr
                node.left = self.grow_tree(X_left, y_left, depth + 1)
                node.right = self.grow_tree(X_right, y_right, depth + 1)
        return node

    def single_predict(self, inputs):
        """对单个样本进行预测"""
        node = self.tree
        while node.left:
            if inputs[node.feature_index] <= node.threshold:
                node = node.left
            else:
                node = node.right
        return node.predicted_class